import glob,sys,subprocess
from multiprocessing import Pool
#from datetime import datetime
from collections import defaultdict

def dealsample(indir: str):
    _id2lib = defaultdict(list)
    files = glob.glob(f'{indir}/*.gz')
    for i in files:
        if i.find('Undetermined') > -1:
            continue
        else:
            sample = i.strip().split('/')[-1].split('.')[0].split('-')[0]
            _id2lib[sample].append(i.strip().split('/')[-1].split('.')[0])
    return _id2lib

ssize = defaultdict(int)
with open(glob.glob(f'{sys.argv[1]}/*.fastq.Stats.xls')[0]) as ds:
    ds.readline()
    for d in ds:
        if d.strip():
            s = d.strip().split()
            ssize[s[0].split('-')[0]] += int(s[2].replace(",",""))
        else:
            break

def fastqc(indir,key,il):
    num = len(il)
    script = ''
    if num == 1:
        if ssize[key] <= 20000000:
            script = f'''export PATH=/data/softwares/fastqc/FastQC:$PATH
            fastqc -q -o {indir} -t 10 {indir}/{key}*.gz
            unzip {indir}/{key}*_fastqc.zip -d {indir}
#            /home/yong_sun/bin/workflow/venv/bin/python /home/yong_sun/bin/workflow/qcstat/q30.py {indir}/{key}*_fastqc/fastqc_data.txt {indir}/{il[0]}.q30
            cp {indir}/{key}*_fastqc/Images/per_base_quality.png {indir}/{key}.per_base_quality.png
            rm -rf {indir}/{key}*_fastqc.zip {indir}/{key}*_fastqc {indir}/{key}*_fastqc.html
            cp {indir}/{key}*.per_base_quality.png /data/mngsSYS/b/reportTMP/1/QC'''
            try:
                a = subprocess.run(script, shell=True, stderr=subprocess.PIPE)
            except Exception as e:
                print(e)
        else:
            script = f'''export PATH=/data/softwares/fastqc/FastQC:$PATH
            gunzip -c {indir}/{key}*.gz | head -n 80000000 > {indir}/{key}_20M
            fastqc -q -o {indir} -t 10 {indir}/{key}_20M
            unzip {indir}/{key}*_fastqc.zip -d {indir}
 #           /home/yong_sun/bin/workflow/venv/bin/python /home/yong_sun/bin/workflow/qcstat/q30.py {indir}/{key}*_fastqc/fastqc_data.txt {indir}/{il[0]}.q30
            cp {indir}/{key}*_fastqc/Images/per_base_quality.png {indir}/{key}.per_base_quality.png
            rm -rf {indir}/{key}*_fastqc.zip {indir}/{key}*_fastqc {indir}/{key}*_fastqc.html {indir}/{key}_20M
            cp {indir}/{key}*.per_base_quality.png /data/mngsSYS/b/reportTMP/1/QC'''
            try:
                a = subprocess.run(script, shell=True, stderr=subprocess.PIPE)
            except Exception as e:
                print(e)
    else:
        if ssize[key] <= 20000000:
            script = f'''export PATH=/data/softwares/fastqc/FastQC:$PATH
            cat {indir}/{key}*.gz > {indir}/{key}_all.fastq.gz
            fastqc -q -o {indir} -t 10 {indir}/{key}_all.fastq.gz
            unzip {indir}/{key}*_fastqc.zip -d {indir}
            cp {indir}/{key}*_fastqc/Images/per_base_quality.png {indir}/{key}.per_base_quality.png
            '''
            try:
                a = subprocess.run(script, shell=True, stderr=subprocess.PIPE)
            except Exception as e:
                print(e)
        else:
            script = f'''export PATH=/data/softwares/fastqc/FastQC:$PATH
            cat {indir}/{key}*.gz | gunzip -c - | head -n 120000000 > {indir}/{key}_30M
            fastqc -q -o {indir} -t 10 {indir}/{key}_30M
            unzip {indir}/{key}*_fastqc.zip -d {indir}
            cp {indir}/{key}*_fastqc/Images/per_base_quality.png {indir}/{key}.per_base_quality.png
            '''
            try:
                a = subprocess.run(script, shell=True, stderr=subprocess.PIPE)
            except Exception as e:
                print(e)
  #      for i in il:
  #          subprocess.run(f'''/home/yong_sun/bin/workflow/venv/bin/python /home/yong_sun/bin/workflow/qcstat/q30.py {indir}/{key}*_fastqc/fastqc_data.txt {indir}/{i}.q30''', shell=True, stderr=subprocess.PIPE)            

        for i in il:
            if i.find('-R') >-1 and i.find('-YF') == -1 and i.find('-SY') == -1: 
                script = f'''export PATH=/data/softwares/fastqc/FastQC:$PATH
                gunzip -c {indir}/{i}*.gz | head -n 80000000 > {indir}/{i}_20M
                fastqc -q -o {indir} -t 10 {indir}/{i}_20M
                unzip {indir}/{i}*_fastqc.zip -d {indir}
                cp {indir}/{i}*_fastqc/Images/per_base_quality.png {indir}/{key}-R.per_base_quality.png
                rm -rf {indir}/{i}*_fastqc.zip {indir}/{i}*_fastqc {indir}/{i}*_fastqc.html {indir}/{i}_20M
                cp {indir}/{key}-R.per_base_quality.png /data/mngsSYS/b/reportTMP/1/QC'''
                try:
                    a = subprocess.run(script, shell=True, stderr=subprocess.PIPE)
                except Exception as e:
                    print(e)

        script =f'''rm -rf {indir}/{key}*_fastqc.zip {indir}/{key}*_fastqc {indir}/{key}*_fastqc.html {indir}/{key}_all.fastq.gz {indir}/{key}_30M
        cp {indir}/{key}*.per_base_quality.png /data/mngsSYS/b/reportTMP/1/QC'''
        try:
            a = subprocess.run(script, shell=True, stderr=subprocess.PIPE)
        except Exception as e:
            print(e)

if __name__ == "__main__":
    p = Pool(30)
    inputdir = sys.argv[1]
    id2lib = dealsample(inputdir)
    res = []
    for k,v in id2lib.items():
        res.append(p.apply_async(fastqc,(inputdir,k,v)))
    p.close()
    p.join()
